// src/transport.rs
use async_trait::async_trait;
use tokio::net::TcpStream;
use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
use tokio::io::{AsyncRead, AsyncWrite, AsyncReadExt, AsyncWriteExt};
use futures_util::{SinkExt, StreamExt};


#[async_trait]
pub trait Transport: Send {
    async fn read(&mut self) -> anyhow::Result<Option<Vec<u8>>>;
    async fn write(&mut self, data: &[u8]) -> anyhow::Result<()>;
    fn is_closed(&self) -> bool;
    #[allow(dead_code)]
    async fn shutdown(&mut self) -> anyhow::Result<()>;
}

pub struct WebSocketTransport<S: AsyncRead + AsyncWrite + Unpin + Send> {
    ws: WebSocketStream<S>,
    closed: bool,
}

impl<S: AsyncRead + AsyncWrite + Unpin + Send> WebSocketTransport<S> {
    pub fn new(ws: WebSocketStream<S>) -> Self {
        Self { ws, closed: false }
    }

    pub async fn send_metadata(&mut self, host: &str, port: u16) -> anyhow::Result<()> {
        let metadata = serde_json::json!({
            "host": host,
            "port": port
        });
        self.ws.send(Message::Text(metadata.to_string())).await?;
        Ok(())
    }
}

#[async_trait]
impl<S: AsyncRead + AsyncWrite + Unpin + Send + 'static> Transport for WebSocketTransport<S> {
    async fn read(&mut self) -> anyhow::Result<Option<Vec<u8>>> {
        match self.ws.next().await {
            Some(Ok(Message::Binary(data))) => Ok(Some(data)),
            Some(Ok(Message::Text(text))) => Ok(Some(text.into_bytes())),
            Some(Ok(Message::Close(_))) => {
                self.closed = true;
                Ok(None)
            }
            Some(Ok(_)) => Ok(None),
            Some(Err(e)) => Err(e.into()),
            None => {
                self.closed = true;
                Ok(None)
            }
        }
    }

    async fn write(&mut self, data: &[u8]) -> anyhow::Result<()> {
        self.ws.send(Message::Binary(data.to_vec())).await?;
        Ok(())
    }

    fn is_closed(&self) -> bool {
        self.closed
    }

    async fn shutdown(&mut self) -> anyhow::Result<()> {
        self.ws.close(None).await?;
        Ok(())
    }
}

pub struct TcpTransport {
    tcp: TcpStream,
    closed: bool,
}

impl TcpTransport {
    pub fn new(tcp: TcpStream) -> Self {
        Self { tcp, closed: false }
    }
}

#[async_trait]
impl Transport for TcpTransport {
    async fn read(&mut self) -> anyhow::Result<Option<Vec<u8>>> {
        let mut buffer = [0u8; 8192];
        match self.tcp.read(&mut buffer).await {
            Ok(0) => {
                self.closed = true;
                Ok(None)
            }
            Ok(n) => Ok(Some(buffer[..n].to_vec())),
            Err(e) => Err(e.into()),
        }
    }

    async fn write(&mut self, data: &[u8]) -> anyhow::Result<()> {
        self.tcp.write_all(data).await?;
        Ok(())
    }

    fn is_closed(&self) -> bool {
        self.closed
    }

    async fn shutdown(&mut self) -> anyhow::Result<()> {
        self.tcp.shutdown().await?;
        Ok(())
    }
}


#[cfg(test)]
mod tests {
    use super::*;
    use tokio::net::TcpListener;
    use tokio_tungstenite::{accept_async, connect_async};

    #[tokio::test]
    async fn test_tcp_transport_echo() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let mut transport = TcpTransport::new(stream);
            let data = transport.read().await.unwrap().unwrap();
            transport.write(&data).await.unwrap();
        });

        let client_handle = tokio::spawn(async move {
            let stream = TcpStream::connect(addr).await.unwrap();
            let mut transport = TcpTransport::new(stream);
            transport.write(b"hello").await.unwrap();
            transport.read().await.unwrap().unwrap()
        });

        let response = client_handle.await.unwrap();
        assert_eq!(response, b"hello");
        server_handle.await.unwrap();
    }

    #[tokio::test]
    async fn test_websocket_transport_echo() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let ws_stream = accept_async(stream).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Read metadata
            transport.read().await.unwrap();

            let data = transport.read().await.unwrap().unwrap();
            transport.write(&data).await.unwrap();
        });

        let client_handle = tokio::spawn(async move {
            let (ws_stream, _) = connect_async(format!("ws://{}", addr)).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            transport.send_metadata("example.com", 80).await.unwrap();
            transport.write(b"world").await.unwrap();
            transport.read().await.unwrap().unwrap()
        });

        let response = client_handle.await.unwrap();
        assert_eq!(response, b"world");
        server_handle.await.unwrap();
    }
}